#include "sms_publisher.h"
#include <iostream>
// #include <sys/mman.h>
// #include <fcntl.h>
// #include <unistd.h>
#include <regex>
#include <random>


#define BUFF_SIZE 1024 * 16


namespace sms {


Publisher::Publisher
(
    std::string topic_url,
    std::string topic_type,
    QoS qos,
    std::string ip,
    int port
)
{
    this->_topic_url = topic_url;
    this->_topic_type = topic_type;

    this->_ip = ip;
    this->_port = port;

    nlohmann::json all_types = load_msg_types();
    if (!all_types.contains(topic_type))
    {
        std::cerr << "[SMS_ERROR]: (" << topic_type << ") " << ec2str(205) << std::endl;
        std::abort();
    }
    int url_err = check_topic_url(topic_url);
    if (url_err != 0)
    {
        std::cerr << "[SMS_ERROR]: (" << topic_url << ") " << ec2str(url_err) << std::endl;
        std::abort();
    }

    this->_last_send_time = 0.0;
    this->_last_upload_time = 0.0;
    this->_force_quit = false;
    this->_suspended = false;
    this->_heartbeat_running = false;
    this->_running = true;
    this->_qos = qos;
    
    this->_client_socket = -1;
    this->_buf = new char[BUFF_SIZE + 1];
    this->_recv_t = nullptr;
    this->_send_t = nullptr;

    this->_error_cnt = 0;
    this->_upload_id = 0;
    this->_uploaded_times.resize(65536, 0);
    this->_transmission_delay = 0.0;

    this->_use_shm = false;
    // this->_shm_fd = -1;
    // this->_shm_size = 1024 * 1024 * 60;
    // this->_shm_ptr = nullptr;
    // this->_shm_offset = 0;
    if (_startswith(topic_type, "memory_msgs::"))
    {
        std::regex pattern("/");
        this->_shm_name = "/" + std::regex_replace(topic_url, pattern, "_");

        init_pubmem(this->_shm_name);
        this->_use_shm = true;

        /*
        this->_shm_fd = shm_open(this->_shm_name.c_str(), O_CREAT | O_RDWR, 0666);
        if (this->_shm_fd == -1)
        {
            std::cerr << "Failed to create shared memory object" << std::endl;
            std::abort();
        }
        if (ftruncate(this->_shm_fd, this->_shm_size) == -1)
        {
            std::cerr << "Failed to set shared memory size" << std::endl;
            std::abort();
        }
        this->_use_shm = true;
        this->_shm_ptr = (char*)mmap(0, this->_shm_size, PROT_READ | PROT_WRITE, MAP_SHARED, this->_shm_fd, 0);
        if (this->_shm_ptr == MAP_FAILED)
        {
            std::cerr << "Failed to map shared memory" << std::endl;
            std::abort();
        }
        */
    }

    this->_recv_t = new std::thread(&Publisher::recv_loop, this);
    this->_send_t = new std::thread(&Publisher::send_loop, this);
    this->_heartbeat_t = new std::thread(&Publisher::heartbeat_loop, this);

    this->_link();
}

Publisher::~Publisher()
{
    this->_force_quit = true;
    this->join();
    if (this->_buf)
        delete this->_buf;
    if (this->_recv_t)
        delete this->_recv_t;
    if (this->_send_t)
        delete this->_send_t;
    if (this->_heartbeat_t)
        delete this->_heartbeat_t;
    
    /*
    if (this->_use_shm)
    {
        if (this->_shm_ptr != nullptr)
        {
            munmap(this->_shm_ptr, this->_shm_size);
            this->_shm_ptr = nullptr;
        }
        close(this->_shm_fd);
        shm_unlink(this->_shm_name.c_str());
        this->_use_shm = false;
    }
    */
}

void Publisher::join()
{
    if (this->_recv_t)
        this->_recv_t->join();
    if (this->_send_t)
        this->_send_t->join();
    if (this->_heartbeat_t)
        this->_heartbeat_t->join();
}

void Publisher::kill()
{
    if (!this->_force_quit)
    {
        {
            std::unique_lock<std::mutex> lock(this->_queue_mtx);
            this->_send_queue.push({});
        }
        this->_cv.notify_one();
        this->_force_quit = true;
        this->_running = false;
        this->_close_socket();
        /*
        if (this->_use_shm)
        {
            if (this->_shm_ptr != nullptr)
            {
                munmap(this->_shm_ptr, this->_shm_size);
                this->_shm_ptr = nullptr;
            }
            close(this->_shm_fd);
            shm_unlink(this->_shm_name.c_str());
            this->_use_shm = false;
        }
        */
    }
}

bool Publisher::publish(nlohmann::json json_msg)
{
    // this->_enforce_publish = enforce;
    if (!this->_suspended && this->_heartbeat_running 
        && ((QoS::BestEffort == this->_qos && get_time_sec() - this->_last_upload_time > this->_transmission_delay * 0.3) 
        || (QoS::Reliability == this->_qos)))
    {
        if (json_msg["timestamp"] == 0.0)
            json_msg["timestamp"] = get_time_sec();
        nlohmann::json topic_upload = def_msg("_sys_msgs::TopicUpload");
        topic_upload["topic"] = json_msg;
        this->_upload_id += 1;
        if (this->_upload_id >= 65535)
            this->_upload_id = 1;
        topic_upload["id"] = this->_upload_id;

        if ((QoS::BestEffort == this->_qos && this->_send_queue.empty()) || (QoS::Reliability == this->_qos))
        {
            if (this->_send_queue.size() >= 1024)
            {
                std::cout << "[SMS_WARN]: " << ec2str(300) << std::endl;
                return false;
            }
            this->_ids_mtx.lock();
            this->_uploaded_times[this->_upload_id] = get_time_sec();
            this->_ids_mtx.unlock();
            {
                std::unique_lock<std::mutex> lock(this->_queue_mtx);
                this->_send_queue.push(topic_upload);
            }
            this->_cv.notify_one();

            this->_last_send_time = get_time_sec();
            this->_last_upload_time = get_time_sec();
            return true;
        }
    }
    return false;
}


void Publisher::send_loop()
{
    while (this->_running)
    {
        std::unique_lock<std::mutex> lock(this->_queue_mtx);
        this->_cv.wait(lock, [this] { return !this->_send_queue.empty(); });

        nlohmann::json topic_upload = this->_send_queue.front();
        if (!topic_upload.empty())
        {
            std::string bytes = encode_msg(topic_upload);
            this->_send_queue.pop();

            this->_send_mtx.lock();
            ssize_t ret = send(this->_client_socket, bytes.c_str(), bytes.size(), MSG_NOSIGNAL | MSG_DONTWAIT);
            this->_send_mtx.unlock();
        }
    }
}


void Publisher::recv_loop()
{
    while (this->_running)
    {
        // std::cout << "_running" << std::endl;
        if (this->_force_quit)
        {
            break;
        }
        if (this->_client_socket < 0)
        {
            // std::cout << "recv_loop() -> this->_client_socket < 0" << std::endl;
            this->_heartbeat_running = false;
            sleep(1);
            continue;
        }
        ssize_t buf_len = read(this->_client_socket, this->_buf, BUFF_SIZE);
        if (buf_len <= 0)
        {
            // std::cout << "recv_loop() -> buf_len == 0" << std::endl;
            this->_close_socket();
            this->_heartbeat_running = false;
            sleep(1);
            continue;
        }
        this->_buf[buf_len] = 0;
        std::string data(this->_buf, buf_len);
        // std::cout << "buf_len: " << buf_len << std::endl;

        std::vector<std::string> checked_msgs;
        std::vector<std::string> parted_msgs;
        std::vector<int> parted_lens;
        std::vector<std::string> recv_msgs;

        _check_msg(data, checked_msgs, parted_msgs, parted_lens);
        if (parted_msgs.size() > 0)
        {
            for (int i=0; i<parted_msgs.size(); i++)
            {
                if (parted_lens[i] > 0)
                {
                    this->_last_msg = parted_msgs[i];
                    this->_last_msg_len = parted_lens[i];
                }
                else
                {
                    this->_last_msg.append(parted_msgs[i]);
                    if (this->_last_msg_len > 0 && this->_last_msg_len <= this->_last_msg.size())
                    {
                        recv_msgs.push_back(this->_last_msg.substr(0, this->_last_msg_len));
                        this->_last_msg_len = 0;
                        this->_last_msg.clear();
                    }
                }
            }
        }
        for (int i=0; i<checked_msgs.size(); i++)
            recv_msgs.push_back(checked_msgs[i]);

        if (recv_msgs.size() > 0)
        {
            for (int i=0; i<recv_msgs.size(); i++)
                this->_parse_msg(recv_msgs[i]);
        }
    }
}

void Publisher::heartbeat_loop()
{
    std::random_device rd;
    std::mt19937 rng(rd());
    std::uniform_int_distribution<int> dist(1, 3);

    int n_try = 0;
    while (this->_running)
    {
        while (this->_heartbeat_running)
        {
            if (this->_force_quit)
            {
                break;
            }
            if (this->_client_socket < 0)
            {
                // std::cout << "send_loop() -> this->_client_socket < 0" << std::endl;
                this->_heartbeat_running = false;
            }
            else
            {
                if (get_time_sec() - this->_last_send_time >= 2.0)
                {
                    this->_heartbeat();
                }
                this->_delay_packet_loss_rate();
            }
            sleep(dist(rng));
        }
        if (this->_force_quit)
        {
            break;
        }
        sleep(1);
        if (!this->_heartbeat_running)
        {
            n_try++;
            if (n_try > 5)
            {
                n_try = 0;
                // std::cout << "-- strat this->_link(): " << std::endl;
                bool succ = this->_link();
                // std::cout << "-- end   this->_link(): " << succ << std::endl;
            }
        }
    }
}

void Publisher::_heartbeat()
{
    nlohmann::json heartbeat_msg = def_msg("_sys_msgs::Publisher");
    heartbeat_msg["topic_type"] = this->_topic_type;
    heartbeat_msg["url"] = this->_topic_url;
    heartbeat_msg["qos"] = this->_qos;
    std::string bytes = encode_msg(heartbeat_msg);

    this->_send_mtx.lock();
    ssize_t ret = send(this->_client_socket, bytes.c_str(), bytes.size(), MSG_NOSIGNAL | MSG_DONTWAIT);
    this->_send_mtx.unlock();

    this->_last_send_time = get_time_sec();
}

void Publisher::_close_socket()
{
	if (this->_client_socket > 0)
    {
		close(this->_client_socket);
	}
	this->_client_socket = -1;
}

bool Publisher::_link()
{
    if (this->_client_socket != -1)
    {
		this->_close_socket();
	}
	this->_client_socket = socket(PF_INET , SOCK_STREAM , 0);
	if (this->_client_socket == -1)
    {
		return false;
	}

    memset(&this->_server_addr , 0 , sizeof(this->_server_addr));
	this->_server_addr.sin_family = AF_INET;
	this->_server_addr.sin_port = htons(this->_port);
    inet_pton(AF_INET, this->_ip.c_str(), &this->_server_addr.sin_addr);

    struct timeval timeout;
    timeout.tv_sec = 5;
    timeout.tv_usec = 0;
    if (setsockopt(this->_client_socket, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) < 0)
    {
        this->_close_socket();
		return false;
    }

    if (connect(this->_client_socket, (struct sockaddr*)&this->_server_addr, sizeof(this->_server_addr)) == -1)
    {
		this->_close_socket();
		return false;
	}
    
    // this->_ids_mtx.lock();
    // this->_uploaded_times.clear();
    // this->_upload_id = 0;
    // this->_ids_mtx.unlock();
    this->_last_msg_len = 0;
    this->_last_msg.clear();
    this->_heartbeat();
    this->_heartbeat_running = true;
    return true;
}

void Publisher::_parse_msg(std::string msg)
{
    nlohmann::json json_msg;
    if (decode_msg(msg, json_msg))
    {
        if (json_msg["type"] == "_sys_msgs::Suspend")
        {
            this->_suspended = true;
            // std::cout << "_sys_msgs::Suspend" << std::endl;
        }
        else if (json_msg["type"] == "_sys_msgs::Unsuspend")
        {
            this->_suspended = false;
            // std::cout << "_sys_msgs::Unsuspend" << std::endl;
        }
        else if (json_msg["type"] == "_sys_msgs::Result")
        {
            if (json_msg["id"] > 0)
            {
                int recv_id = json_msg["id"];
                this->_error_cnt = 0;

                this->_ids_mtx.lock();
                if (recv_id >= 0 && recv_id < this->_uploaded_times.size())
                {
                    if (this->_uploaded_times[recv_id] > 0.0 && get_time_sec() - this->_uploaded_times[recv_id] > 0.0)
                    {
                        this->_delays.push_back(get_time_sec() - this->_uploaded_times[recv_id]);
                    }
                    else
                    {
                        this->_uploaded_times[recv_id] = 0.0;
                    }
                }
                this->_ids_mtx.unlock();
            }
            if (json_msg["error_code"] > 0)
            {
                std::cout << "[SMS_ERROR]: " << json_msg["data"] << std::endl;
                this->_error_cnt ++;
                if (this->_error_cnt > 5)
                    this->_suspended = true;
            }
        }
        else if (json_msg["type"] != "_sys_msgs::HeartBeat")
        {
            std::cout << "[SMS_ERROR]: " << json_msg.dump() << std::endl;
        }
    }
}


void Publisher::_delay_packet_loss_rate()
{
    double delay = 0.0;
    int delay_cnt = 0;

    this->_ids_mtx.lock();
    for (size_t i=0; i<this->_delays.size(); i++)
    {
        delay += this->_delays[i];
        delay_cnt ++;
    }
    this->_delays.clear();
    this->_ids_mtx.unlock();

    if (delay_cnt > 0)
    {
        delay = delay / delay_cnt;
        this->_transmission_delay = delay;
    }

    // std::cout << "this->_transmission_delay: " << this->_transmission_delay << ", size: " << delay_cnt << std::endl;
}


nlohmann::json Publisher::cvimg2sms_mem(const cv::Mat& cvimg, const std::string frame_id, const double timestamp)
{
    if (this->_use_shm)
    {
        return sms::cvimg2sms_mem(cvimg, this->_shm_name, frame_id, timestamp);
    }
    else
    {
        std::cout << "[SMS_ERROR]: " << ec2str(501) << std::endl;
        return {};
    }
    /*
    if (this->_use_shm && this->_shm_ptr != nullptr)
    {
        nlohmann::json img_msg = def_msg("memory_msgs::RawImage");
        img_msg["timestamp"] = get_time_sec();
        img_msg["frame_id"] = frame_id;
        img_msg["height"] = cvimg.rows;
        img_msg["width"] = cvimg.cols;
        if (cvimg.type() == CV_8UC3)
            img_msg["encoding"] = "8UC3";
        img_msg["shm_name"] = this->_shm_name;
        size_t img_siz = cvimg.total() * cvimg.elemSize();
        if (this->_shm_offset + img_siz > this->_shm_size)
        {
            this->_shm_offset = 0;
        }
        std::memcpy(this->_shm_ptr + this->_shm_offset, cvimg.data, img_siz);
        img_msg["offset"] = this->_shm_offset;
        this->_shm_offset += img_siz;
        return img_msg;
    }
    else
    {
        std::cerr << "Failed to use shared memory in cvimg2shm" << std::endl;
        std::abort();
        return {};
    }
    */
}

nlohmann::json Publisher::pcl2sms_mem(const cv::Mat& pcl, const nlohmann::json fields, const std::string frame_id, const double timestamp)
{
    if (this->_use_shm)
    {
        nlohmann::json msg = sms::pcl2sms_mem(pcl, this->_shm_name, frame_id, timestamp);
        msg["fields"] = fields;
        return msg;
    }
    else
    {
        std::cout << "[SMS_ERROR]: " << ec2str(501) << std::endl;
        return {};
    }
}


Logger::Logger(std::string name, bool to_screen, std::string ip, int port)
{
    this->_name = name;
    this->_to_screen = to_screen;
    this->_pub = new Publisher(
        "/log/" + name,
        "std_msgs::String",
        QoS::Reliability,
        ip,
        port
    );
}

Logger::~Logger()
{

}

void Logger::quit()
{
    this->_pub->kill();
    this->_pub->join();
}

void Logger::info(std::string msg)
{
    std::string f_msg = this->_format_msg(msg, "INFO");
    if (this->_to_screen)
    {
        std::cout << "\033[32m" << f_msg  << "\033[0m" << std::endl;
    }
}

void Logger::warn(std::string msg)
{
    std::string f_msg = this->_format_msg(msg, "WARN");
    if (this->_to_screen)
    {
        std::cout << "\033[33m" << f_msg  << "\033[0m" << std::endl;
    }
}

void Logger::error(std::string msg)
{
    std::string f_msg = this->_format_msg(msg, "ERROR");
    if (this->_to_screen)
    {
        std::cout << "\033[31m" << f_msg  << "\033[0m" << std::endl;
    }
}

void Logger::debug(std::string msg)
{
    std::string f_msg = this->_format_msg(msg, "DEBUG");
    if (this->_to_screen)
    {
        std::cout << f_msg << std::endl;
    }
}

std::string Logger::_format_msg(std::string msg, std::string msg_type)
{
    std::string msg_data = "[" + msg_type + "] " + get_time_str_msec() + " - " + msg;

    nlohmann::json sms_msg = sms::def_msg("std_msgs::String");
    sms_msg["data"] = msg_data;
    this->_pub->publish(sms_msg);
    return msg_data;
}


CSVLogger::CSVLogger(std::vector<std::string> columns, std::string name, std::string ip, int port)
{
    this->_name = name;
    this->_columns = columns;
    this->_pub = new Publisher(
        "/csv_log/" + name,
        "std_msgs::String",
        QoS::Reliability,
        ip,
        port
    );

    nlohmann::json sms_msg = sms::def_msg("std_msgs::String");
    std::string msg_data = "#-";
    for (size_t i=0; i<columns.size(); i++)
    {
        msg_data += columns[i];
        if (i != columns.size() - 1)
            msg_data += ",";
    }
    sms_msg["data"] = msg_data;
    this->_pub->publish(sms_msg);
}

CSVLogger::~CSVLogger()
{

}

void CSVLogger::quit()
{
    this->_pub->kill();
    this->_pub->join();
}

void CSVLogger::append(std::vector<double> values)
{
    assert(values.size() == this->_columns.size());
    nlohmann::json sms_msg = sms::def_msg("std_msgs::String");
    std::string msg_data = "";
    for (size_t i=0; i<values.size(); i++)
    {
        msg_data += std::to_string(values[i]);
        if (i != values.size() - 1)
            msg_data += ",";
    }
    sms_msg["data"] = msg_data;
    this->_pub->publish(sms_msg);
}


}
